#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import logging
from typing import Dict, Any, Optional, List
import uuid
import argparse
import sys

# 添加项目根目录到Python搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastmcp import FastMCP, Context
from pydantic import BaseModel, Field

from vector_store.factory import VectorStoreFactory
from graph.graph_connection import NebulaGraphConnection
from mcp_server.context_retriever import ContextRetriever
from config import config

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 定义查询参数模型
class QueryParams(BaseModel):
    query_text: str = Field(..., description="查询文本，用于检索相关代码上下文")
    top_k: int = Field(3, description="每种节点类型返回的最大结果数")

# 创建FastMCP实例
mcp = FastMCP("代码上下文检索服务器", dependencies=["pydantic", "logging"])

# 初始化服务器组件
vector_store = None
graph_connection = None
context_retriever = None

def initialize_components():
    """初始化服务器组件"""
    global vector_store, graph_connection, context_retriever
    
    # 初始化向量存储
    vector_store = VectorStoreFactory.create_vector_store()
    
    # 初始化图数据库连接
    nebula_config = config.get_nebula_config()
    graph_connection = NebulaGraphConnection(nebula_config)
    connected = graph_connection.connect()
    if not connected:
        logger.error("无法连接到NebulaGraph数据库")
        raise ConnectionError("无法连接到NebulaGraph数据库")
    
    # 初始化上下文检索器
    context_retriever = ContextRetriever(
        vector_store=vector_store,
        graph_connection=graph_connection
    )
    
    logger.info("代码上下文MCP服务器初始化完成")

@mcp.tool()
async def start_session() -> str:
    """
    你必须在检索上下文之前使用该工具开启新的对话
    """
    # 生成唯一的会话ID
    session_id = str(uuid.uuid4())
    logger.info(f"创建新会话: {session_id}")
    return session_id

@mcp.tool()
async def retrieve_code_context(sessionId:str, original_user_query_text:str,query_text: str, top_k: int = 3, ctx: Context = None, mcp_resort: bool=False) -> Dict[str, Any]:
    """
    根据查询文本检索相关代码上下文
    
    Args:
        sessionId：当前会话的sessionId
        original_user_query_text: 用户原始查询文本
        query_text: 你想要查询的文本
        top_k: 每种节点类型返回的最大结果数
        ctx: FastMCP上下文对象
        mcp_resort: 客户端必须明确声明支持MCP resort，才赋值true
    Returns:
        包含相关上下文的字典
    """
    # 当前现状：
    # 1.根据用户提问找到对应的向量节点
    # 2.拿到向量节点在知识图谱的关联节点
    # 理想化：无需重复向量化
    # 1.retrieve_code_context 查到问题的关联节点
    # 2.retrieve_code_graph 入参:图id 过程:图数据库里查关联 返回:关联节点
    # todo
    # 1.查找到上下文 √
    # 2.上下文保存到当前session中
    # 3.如果上下文超出限制则进行重排，并返回resort标记
    try:

        if ctx:
            ctx.info(f"接收到上下文检索请求: {query_text}")
        
        result = context_retriever.retrieve_context(query_text, top_k)
        return result
    except Exception as e:
        error_msg = f"检索上下文时发生错误: {str(e)}"
        if ctx:
            ctx.error(error_msg)
        return {
            "error": str(e),
            "query": query_text,
            "nodes": {},
            "total_nodes": 0
        }

@mcp.prompt()
def code_context_prompt(query: str) -> str:
    """创建代码上下文检索的提示模板"""
    return f"""请根据以下查询检索相关代码上下文：
查询: {query}

请提供最相关的代码片段和上下文信息。"""

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='代码上下文MCP服务器')
    parser.add_argument('--host', type=str, default='127.0.0.1', help='服务器主机地址')
    parser.add_argument('--port', type=int, default=8000, help='服务器端口')
    parser.add_argument('--vector-db', type=str, help='向量数据库路径')
    parser.add_argument('--nebula-ip', type=str, help='NebulaGraph IP地址')
    parser.add_argument('--nebula-port', type=int, help='NebulaGraph端口')
    parser.add_argument('--nebula-user', type=str, help='NebulaGraph用户名')
    parser.add_argument('--nebula-password', type=str, help='NebulaGraph密码')
    parser.add_argument('--nebula-space', type=str, help='NebulaGraph工作空间')
    parser.add_argument('--api-key', type=str, help='嵌入API密钥')
    parser.add_argument('--log-level', type=str, choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], 
                      default='INFO', help='日志级别')
    return parser.parse_args()

def main():
    """主函数"""
    args = parse_args()
    
    # 更新配置
    if args.vector_db:
        os.environ["SQLITE_DB_PATH"] = args.vector_db
    
    if args.nebula_ip:
        os.environ["NEBULA_IP"] = args.nebula_ip
    if args.nebula_port:
        os.environ["NEBULA_PORT"] = str(args.nebula_port)
    if args.nebula_user:
        os.environ["NEBULA_USER"] = args.nebula_user
    if args.nebula_password:
        os.environ["NEBULA_PASSWORD"] = args.nebula_password
    if args.nebula_space:
        os.environ["NEBULA_SPACE"] = args.nebula_space
        
    if args.api_key:
        os.environ["EMBEDDING_API_KEY"] = args.api_key
    
    # 设置日志级别
    logging.getLogger().setLevel(getattr(logging, args.log_level))
    
    # 更新配置
    global config
    from config.config import Config
    config = Config()
    
    # 初始化组件
    initialize_components()
    
    # 启动服务器
    mcp.run(host=args.host, port=args.port)

# 初始化组件
initialize_components()

if __name__ == "__main__":
    main()
